#include <iostream>
#include <vector>
#include <cassert>
#include "IndexMinHeap.h"

using namespace std;


template<typename Graph, typename Weight>
class PrimMST{

private:
    Graph &G;						// 存储图的引用
    vector<Edge<Weight>> mst;		// 最小生成树所有的边

    bool* marked;					// 点是否被遍历过
    IndexMinHeap<Weight> ipq;		// 最小索引堆
    vector<Edge<Weight>*> edgeTo;	// 和每个节点相连的最短横切边
    Weight mstWeight;				// 最小生成树的最小权值


    void visit(int v){
        assert( !marked[v] );
        marked[v] = true;			// 访问该顶点将其标记

        typename Graph::adjIterator adj(G,v);
        for( Edge<Weight>* e = adj.begin() ; !adj.end() ; e = adj.next() ){
			// v 相邻的顶点 w
            int w = e->other(v);
            if( !marked[w] ){

                if( !edgeTo[w] ){	// 没有找到过和w相邻的横切边
                    edgeTo[w] = e;
                    ipq.insert(w, e->wt()); // 将这条边加入最小索引堆
                }
				// 之前找到过和w，相邻的横切边,比较e新找到边的权值是否小于
				// 之前找到和w相邻横切边的权值
				else if( e->wt() < edgeTo[w]->wt() ){

                    edgeTo[w] = e;			// 扔掉之前找到的邻边
                    ipq.change(w, e->wt());	// 维护索引堆权值
                }
            }
			// 被标记过不做任何处理
        }

    }
public:
    // assume graph is connected
    PrimMST(Graph &graph):G(graph), ipq(IndexMinHeap<double>(graph.V())){

        assert( graph.E() >= 1 );

        marked = new bool[G.V()];
        for( int i = 0 ; i < G.V() ; i ++ ){
            marked[i] = false;
            edgeTo.push_back(NULL);
        }
        mst.clear();

        visit(0);
		// 最小索引堆不为空,还需要考虑横切边
        while( !ipq.isEmpty() ){
            int v = ipq.extractMinIndex();	// 取出最小的横切边的索引
            assert( edgeTo[v] );
            mst.push_back( *edgeTo[v] );	// 将这条边放进最小生成树中
            visit( v );
        }

		// 求出最小权值
        mstWeight = mst[0].wt();
        for( int i = 1 ; i < mst.size() ; i ++ )
            mstWeight += mst[i].wt();
    }

    ~PrimMST(){
        delete[] marked;
    }

    vector<Edge<Weight>> mstEdges(){
        return mst;
    };

    Weight result(){
        return mstWeight;
    };
};
